/**
 * @file linspaced_custom_tiling.cpp
 *
 * Copyright (C) 2025. Huawei Technologies Co., Ltd. All rights reserved.
 *
 * This program is distributed in the hope that it will be useful,
 * but WITHOUT ANY WARRANTY; without even the implied warranty of
 * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.
 */
#include <cassert>
#include <cstddef>
#include "linspaced_custom_tiling.h"

constexpr uint32_t DATA_TYPE_SIZE[] = {2, 4, 2, 4};                                      // 数据类型索引
constexpr uint32_t BLOCK_SIZE = 32;                                                      // 32B内存对齐基准
constexpr uint32_t BUFFER_NUM = 2;                                                       // 双缓冲配置
constexpr uint32_t UB_BLOCK_NUM = 100;                                                   // UB最大可使用的32B块数量
constexpr uint32_t MAX_AVAILABLE_UB_BLOCK_NUM = UB_BLOCK_NUM / BUFFER_NUM * BUFFER_NUM;  // 双缓冲适配的UB块数

void TilingParamsCalc(uint32_t length, uint32_t alignNum,
                      uint32_t &tileNum, uint32_t &tileLength, uint32_t &lastTileLength)
{
    assert(alignNum != 0U);
    uint32_t maxPerCoreElem = alignNum * MAX_AVAILABLE_UB_BLOCK_NUM;  
    tileNum = length / maxPerCoreElem;  

    if ((length % maxPerCoreElem == 0U) || (tileNum == 0U)) 
    {
        if (tileNum == 0U) tileNum = (length + alignNum - 1) / alignNum;                                       // length > alignNum时，
                                                                                                               // 按alignNum拆分Tile
        if (length < maxPerCoreElem) 
        {
            uint32_t blockCount = (length + alignNum - 1) / alignNum;  
            blockCount = (blockCount + BUFFER_NUM - 1) / BUFFER_NUM * BUFFER_NUM;                              // 双缓冲对齐
            tileLength = blockCount * alignNum; 
            lastTileLength = (tileNum == 1) ? length : (length % alignNum == 0 ? alignNum : length % alignNum);// lastTileLength=实际剩余元素数
        } 
        else 
        {
            tileLength = maxPerCoreElem;
            lastTileLength = length - (tileNum - 1) * tileLength;
        }
    } 
    else 
    {
        tileNum++;
        tileLength = maxPerCoreElem;
        lastTileLength = length - (tileNum - 1) * tileLength;
    }
}

void GenerateTilingData(uint8_t *tilingBuf, uint32_t blockDim)
{
    LinspacedCustomTilingData *tiling = reinterpret_cast<LinspacedCustomTilingData *>(tilingBuf);
    uint32_t totalLength = tiling->totalLength; 
    uint32_t dataTypeSize = DATA_TYPE_SIZE[SIG_FLOAT]; //输出float类型
    uint32_t alignNum = BLOCK_SIZE / dataTypeSize;
    uint32_t totalLengthAligned = (totalLength % alignNum == 0U)
        ? totalLength
        : ((totalLength + alignNum - 1) / alignNum) * alignNum;

    if ((totalLengthAligned / alignNum) % blockDim == 0U) 
    {   // 核间均分
        uint32_t blockLength = totalLengthAligned / blockDim;
        uint32_t tileNum, tileLength, lastTileLength;
        TilingParamsCalc(blockLength, alignNum, 
                         tileNum, tileLength, lastTileLength);

        tiling->blockLength = blockLength;
        tiling->tileNum = tileNum;
        tiling->tileLength = tileLength;
        tiling->lastTileLength = lastTileLength;
        tiling->isEvenCore = 1U;
    } 
    else 
    {   // 核间非均分
        uint32_t totalBlockCount = totalLengthAligned / alignNum;  // 总32B块数
        uint32_t formerNum = totalBlockCount % blockDim;           // 多处理1块的核心数
        uint32_t tailNum = blockDim - formerNum;                   // 正常处理的核心数

        uint32_t formerBlockCount = (totalBlockCount + blockDim - 1) / blockDim;
        uint32_t tailBlockCount = totalBlockCount / blockDim;
        uint32_t formerLength = formerBlockCount * alignNum;
        uint32_t tailLength = tailBlockCount * alignNum;

        uint32_t fTileNum, fTileLen, fLastTileLen;
        TilingParamsCalc(formerLength, alignNum, 
                         fTileNum, fTileLen, fLastTileLen);
        
        uint32_t tTileNum, tTileLen, tLastTileLen;
        TilingParamsCalc(tailLength, alignNum, 
                         tTileNum, tTileLen, tLastTileLen);

        tiling->formerNum = formerNum;
        tiling->formerLength = formerLength;
        tiling->formerTileNum = fTileNum;
        tiling->formerTileLength = fTileLen;
        tiling->formerLastTileLength = fLastTileLen;

        tiling->tailNum = tailNum;
        tiling->tailLength = tailLength;
        tiling->tailTileNum = tTileNum;
        tiling->tailTileLength = tTileLen;
        tiling->tailLastTileLength = tLastTileLen;
        tiling->isEvenCore = 0U;
    }
}
